import os
import matplotlib.pyplot as plt
from torchvision import transforms
from torch.utils.data import Dataset
from PIL import Image
from glob import glob

# 假设 utils 模块中有 PIL_Image_ToTensor 函数，如果没有，可以使用 transforms.ToTensor
from utils import PIL_Image_ToTensor


class Neu_Seg_Competition_Dataset(Dataset):
    def __init__(self, data_path, data_type):
        self.data_path = data_path
        self.data_type = data_type
        self.to_tensor = transforms.ToTensor()

        # 获取所有图像的文件路径
        image_paths = glob(os.path.join(data_path, 'images', data_type, '*.jpg'))
        self.image_names = [os.path.splitext(os.path.basename(path))[0] for path in image_paths]

        # 预加载所有图像和掩码到内存
        self.images = []
        self.masks = []
        for name in self.image_names:
            # 加载图像
            image_path = os.path.join(data_path, 'images', data_type, f'{name}.jpg')
            image = Image.open(image_path).convert('RGB')
            image_tensor = self.to_tensor(image)
            self.images.append(image_tensor)

            # 加载掩码
            mask_path = os.path.join(data_path, 'annotations', data_type, f'{name}.png')
            mask = Image.open(mask_path).convert('L')  # 根据实际情况选择转换模式
            mask_tensor = PIL_Image_ToTensor(mask)  # 或者使用 transforms.ToTensor()(mask)
            self.masks.append(mask_tensor)

    def __len__(self):
        return len(self.image_names)

    def __getitem__(self, idx):
        image = self.images[idx]
        mask = self.masks[idx]
        image_name = self.image_names[idx]
        return image, mask, image_name


if __name__ == '__main__':
    import torch

    # 创建数据集实例
    train_dataset = Neu_Seg_Competition_Dataset(data_path='.', data_type='training')
    print(f'Dataset length: {len(train_dataset)}')

    # 获取第121个样本
    index = 120
    image_tensor, mask_tensor, image_name = train_dataset[index]

    # 打印掩码中的唯一值
    unique_values = torch.unique(mask_tensor)
    print(f'Unique values in mask: {unique_values}')

    # 将张量转换为 NumPy 数组以供 matplotlib 使用
    # 对于图像，转换为 HxWxC 格式
    image = image_tensor.permute(1, 2, 0).numpy()
    # 对于掩码，移除通道维度
    mask = mask_tensor.squeeze().numpy()

    # 创建一个包含两张子图的图形
    fig, axs = plt.subplots(1, 2, figsize=(10, 5))  # 1 行 2 列

    # 显示图像
    axs[0].imshow(image)
    axs[0].set_title(f'Image: {image_name}')
    axs[0].axis('off')  # 隐藏坐标轴

    # 显示掩码
    axs[1].imshow(mask, cmap='gray')
    axs[1].set_title('Mask')
    axs[1].axis('off')  # 隐藏坐标轴

    plt.tight_layout()
    plt.show()
